import unittest
import numpy as np
from collections import defaultdict
from dv_hop_positioning import DV_Hop_Positioning, INF


class TestDVHopPositioning(unittest.TestCase):

    def setUp(self):
        # Mock data for testing
        self.anchor_data = {
            "A": np.array([0, 0]),
            "B": np.array([0, 6]),
            "C": np.array([8, 0]),
        }
        self.unknown_data = {"U1": np.array([4, 3]), "U2": np.array([8, 6])}
        self.graph = {
            "A": ["B", "C"],
            "B": ["A", "U1", "U2"],
            "C": ["A", "U1", "U2"],
            "U1": ["C", "B"],
            "U2": ["B", "C"],
        }

        class MockPointsData:
            def __init__(self, anchor_data, unknown_data):
                self.anchor_data = anchor_data
                self.unknown_data = unknown_data

            def anchor_data_dict(self):
                return self.anchor_data

            def unknown_data_dict(self):
                return self.unknown_data

            def get_all_points_name(self):
                return list(self.anchor_data.keys()) + list(self.unknown_data.keys())

        class MockGraph:
            def __init__(self, graph):
                self.graph = graph

            def get_neighbors(self, node):
                return self.graph.get(node, [])

        self.points_data = MockPointsData(self.anchor_data, self.unknown_data)
        self.graph_data = MockGraph(self.graph)

    def test_generate_hop_graph(self):
        dv_hop = DV_Hop_Positioning(self.points_data, self.graph_data)
        expected_hop_graph = defaultdict(
            lambda: defaultdict(lambda: INF),
            {
                "A": {"A": 0, "B": 1, "C": 1, "U1": 2, "U2": 2},
                "B": {"A": 1, "B": 0, "C": 2, "U1": 1, "U2": 1},
                "C": {"A": 1, "B": 2, "C": 0, "U1": 1, "U2": 1},
                "U1": {"A": 2, "B": 1, "C": 1, "U1": 0, "U2": 2},
                "U2": {"A": 2, "B": 1, "C": 1, "U1": 2, "U2": 0},
            },
        )
        self.assertEqual(dv_hop.hop_graph, expected_hop_graph)

    def test_init_dis_graph(self):
        dv_hop = DV_Hop_Positioning(self.points_data, self.graph_data)
        expected_dis_graph = defaultdict(
            lambda: defaultdict(lambda: INF),
            {
                "A": {"A": 0, "B": 6, "C": 8},
                "B": {"A": 6, "B": 0, "C": 10},
                "C": {"A": 8, "B": 10, "C": 0},
            },
        )
        self.assertEqual(dv_hop.dis_graph, expected_dis_graph)

    def test_generate_all_average_dis(self):
        dv_hop = DV_Hop_Positioning(self.points_data, self.graph_data)
        expected_all_average_dis = defaultdict(
            lambda: INF, {"A": 14 / 2, "B": 16 / 3, "C": 18 / 3}
        )
        self.assertEqual(dv_hop.all_average_dis, expected_all_average_dis)

    def test_predict_archor_points_dis(self):
        dv_hop = DV_Hop_Positioning(self.points_data, self.graph_data)
        dv_hop.predict_archor_points_dis()
        expected_dis_graph = defaultdict(
            lambda: defaultdict(lambda: INF),
            {
                "A": {"A": 0, "B": 6, "C": 8, "U1": 14, "U2": 14},
                "B": {"A": 6, "B": 0, "C": 10, "U1": 16 / 3, "U2": 16 / 3},
                "C": {"A": 8, "B": 10, "C": 0, "U1": 18 / 3, "U2": 18 / 3},
            },
        )
        self.assertEqual(dv_hop.dis_graph, expected_dis_graph)

    def test_predict_unknown_points_dis(self):
        dv_hop = DV_Hop_Positioning(self.points_data, self.graph_data)
        dv_hop.predict_unknown_points_dis()
        expected_dis_graph = defaultdict(
            lambda: defaultdict(lambda: INF),
            {
                "A": {"A": 0, "B": 6, "C": 8},
                "B": {"A": 6, "B": 0, "C": 10},
                "C": {"A": 8, "B": 10, "C": 0},
                "U1": {
                    "A": 11.333333333333332,
                    "B": 5.666666666666666,
                    "C": 5.666666666666666,
                },
                "U2": {
                    "A": 11.333333333333332,
                    "B": 5.666666666666666,
                    "C": 5.666666666666666,
                },
            },
        )
        print(dv_hop.dis_graph)
        self.assertEqual(dv_hop.dis_graph, expected_dis_graph)

    # def test_run(self):
    #     dv_hop = DV_Hop_Positioning(self.points_data, self.graph_data)
    #     result = dv_hop.run()
    #     expected_result = {"U1": np.array([5, 5]), "U2": np.array([5, 10])}
    #     for key in expected_result:
    #         np.testing.assert_array_almost_equal(result[key], expected_result[key])


if __name__ == "__main__":
    unittest.main()
